package com.extra;

import java.util.*;

public class Solution1382 {
    List<Integer> list = new ArrayList<>();
    public TreeNode balanceBST(TreeNode root) {
        if (root == null) return null;
        getArray(root);
        int[] nums = new int[list.size()];
        for (int i = 0; i < list.size(); i++) {
            nums[i] = list.get(i);
        }
        return getBalanceTree(0, nums.length, nums);
    }

    public void getArray(TreeNode root) {
        if (root.left != null) {
            getArray(root.left);
        }
        list.add(root.val);
        if (root.right != null) {
            getArray(root.right);
        }
    }

    public TreeNode getBalanceTree(int i, int j, int[] nums) {
        if (i >= j) return null;
        if ((j - i) == 1) return new TreeNode(nums[i]);
        int mid = (i + j) >> 1;
        TreeNode root = new TreeNode(nums[mid]);
        root.left = getBalanceTree(i, mid, nums);
        root.right = getBalanceTree(mid + 1, j, nums);
        return root;
    }
}
